module calculation_mod
  
  use constants_mod
  use IO_mod

  implicit none

contains

subroutine calc_DCT(u, v, dct_u, dct_v)

  real(r8), intent(in) :: u(1:nlons, 1:nlats), v(1:nlons, 1:nlats)
  real(r8), intent(inout) :: dct_u(0:nlons-1, 0:nlats-1), &
                             dct_v(0:nlons-1, 0:nlats-1)
  real(r8), allocatable :: betam(:), betan(:) 
  real(r8) tmp
  integer i, j, m, n, it, ilev

  print*, "Calculating DCT ..."
  allocate(betam(0:nlons-1), betan(0:nlats-1))
  betam(0) = sqrt(1.0 / nlons)
  betam(1:nlons-1) = sqrt(2.0 / nlons)
  betan(0) = sqrt(1.0 / nlats)
  betan(1:nlats-1) = sqrt(2.0 / nlats)

  do n = 0, nlats - 1
    do m = 0, nlons - 1
      DCT_u(m,n) = 0.0
      DCT_v(m,n) = 0.0
      do j = 0, nlats - 1
        do i = 0, nlons - 1
          tmp = cos(pi * m * (i + 0.5) / nlons) * &
                cos(pi * n * (j + 0.5) / nlats)
          DCT_u(m,n) = DCT_u(m,n) + u(i+1,j+1) * tmp 
          DCT_v(m,n) = DCT_v(m,n) + v(i+1,j+1) * tmp
        end do
      end do
      DCT_u(m,n) = betam(m) * betan(n) * DCT_u(m,n)
      DCT_v(m,n) = betam(m) * betan(n) * DCT_v(m,n)
    end do
  end do

  deallocate(betam, betan)  
    
end subroutine calc_DCT

subroutine calc_power_spectral(dct1, dct2, power1, power2, wavelength, wavenumber)
  
  real(r8), intent(in) :: dct1(0:nlons-1, 0:nlats-1), &
                          dct2(0:nlons-1, 0:nlats-1)
  real(r8), intent(inout) :: power1(Nwave), power2(Nwave), &
                            wavelength(Nwave), wavenumber(Nwave)
  real(r8) mu_lower, mu_upper, mu, khat, nlonlat
  integer k, m, n, it, ilev

  print*, "Calculating power spectral ..."
  nlonlat = nlons * nlats

  do k = 1, Nwave
    power1(k) = 0.0
    power2(k) = 0.0
    mu_lower = real(k) / min(nlons, nlats)
    mu_upper = real(k + 1) / min(nlons, nlats)
    wavelength(k) = 2.0 * (radius * dlon * 1.e-03) * min(nlons, nlats) / k !km
    wavenumber(k) = pi * mu_lower / (radius * dlon)
    do n = 0, nlats - 1
      do m = 0, nlons - 1
        if (m == 0 .and. n == 0) cycle
        mu = sqrt(real(m)**2 / real(nlons)**2 + real(n)**2 / real(nlats)**2)
        khat = pi * mu / (radius * dlon)
        if (mu >= mu_lower .and. mu < mu_upper) then
          ! based on divergence and vorticity
          if (vor_div_based) then
            power1(k) = power1(k) + 0.5 * dct1(m,n)**2 / nlonlat / khat**2
            power2(k) = power2(k) + 0.5 * dct2(m,n)**2 / nlonlat / khat**2
          else
            ! based on u and v
            power1(k) = power1(k) + 0.5 * dct1(m,n)**2 / nlonlat * (radius * dlon) * min(nlons, nlats) / pi
            power2(k) = power2(k) + 0.5 * dct2(m,n)**2 / nlonlat * (radius * dlon) * min(nlons, nlats) / pi
          end if
        end if
      end do
    end do
  end do


end subroutine calc_power_spectral

end module calculation_mod  